from collections.abc import Iterable
from typing import Callable, TypeVar

from zkl_aiutils_datasets.basics import Dataset
from zkl_aiutils_datasets.processing import ChainedDatasets, DelegatingDataset
from zkl_aiutils_datasets.wrapping import wrap_dataset
from .threaded_mapped import StateSafeThreadedMappedDataset, ThreadedMappedDataset

SrcSample = TypeVar('SrcSample')
DstSample = TypeVar('DstSample')


class ThreadedAugmentedDataset(DelegatingDataset[DstSample]):
    def __init__(self,
        dataset: Dataset[SrcSample],
        augment_func: Callable[[SrcSample], Iterable[DstSample]], *,
        augment_multiple: int | None = None,
    ):
        self._dataset = dataset
        self._augment_func = augment_func
        self._augment_multiple = augment_multiple
        super().__init__(ChainedDatasets(ThreadedMappedDataset(dataset,
            lambda src_sample: wrap_dataset(augment_func(src_sample), samples_n=augment_multiple))))

    @property
    def dataset(self) -> Dataset[SrcSample]:
        return self._dataset

    @property
    def augment_func(self) -> Callable[[SrcSample], Iterable[DstSample]]:
        return self.augment_func

    @property
    def augment_multiple(self) -> int | None:
        return self._augment_multiple


class StateSafeThreadedAugmentedDataset(DelegatingDataset[DstSample]):
    def __init__(self,
        dataset: Dataset[SrcSample],
        augment_func: Callable[[SrcSample], Iterable[DstSample]], *,
        augment_multiple: int | None = None,
    ):
        self._dataset = dataset
        self._augment_func = augment_func
        self._augment_multiple = augment_multiple
        super().__init__(ChainedDatasets(StateSafeThreadedMappedDataset(dataset,
            lambda src_sample: wrap_dataset(augment_func(src_sample), samples_n=augment_multiple))))

    @property
    def dataset(self) -> Dataset[SrcSample]:
        return self._dataset

    @property
    def augment_func(self) -> Callable[[SrcSample], Iterable[DstSample]]:
        return self.augment_func

    @property
    def augment_multiple(self) -> int | None:
        return self._augment_multiple
